# coding: utf-8

# -------------------------------------------------------------------------------
# Name:         mysql_client.py
# Description:  Mysql客户端
# Author:       XiangjunZhao
# EMAIL:        2419352654@qq.com
# Date:         2020/3/5 13:26
# -------------------------------------------------------------------------------
import logging

import pymysql
from DBUtils.PooledDB import PooledDB
from sshtunnel import SSHTunnelForwarder

logger = logging.getLogger(__name__)


class MysqlClient(object):
    '''
    Mysql客户端
    '''

    def __init__(self, host=None, port=None, user=None, password=None, database=None, charset='utf8'):
        '''
        初始化Mysql客户端
        :param host:主机地址
        :param port:端口
        :param user:用户名
        :param password:密码
        :param database:数据库
        :param charset: 数据库编码
        '''
        # 数据库地址
        self.host = host
        # 数据库端口
        self.port = port or 3306
        # 数据库用户名
        self.user = user
        # 数据库密码
        self.password = password
        # 数据库名称
        self.database = database
        self.charset = charset
        self.db_pool = PooledDB(pymysql, mincached=1, maxcached=5, host=self.host,
                                port=self.port, user=self.user, passwd=self.password,
                                db=self.database, charset=self.charset, connect_timeout=60)

    def get_connet(self):
        '''
        获取数据库连接
        :return:
        '''
        try:
            # 获取数据库连接
            return self.db_pool.connection()
        except BaseException as e:
            logger.warning("获取数据库连接失败，异常：{}".format(e))
            self.get_connet()

    def close_connect(self, conn):
        '''
        关闭数据库连接
        :param conn:数据库连接
        :return:
        '''
        if conn:
            conn.close()

    def close(self):
        '''
        关闭数据库连接池
        @return:
        '''
        try:
            self.db_pool.close()
        except Exception as e:
            logger.error('Mysql数据库断开连接异常，原因：{}'.format(str(e)))

    def query_one(self, conn, sql):
        '''
        查询一条数据
        :param conn:数据库连接
        :param sql:查询sql
        :return:dict
        '''

        # 获取数据库游标
        with conn.cursor() as cursor:
            # 执行查询sql语句
            cursor.execute(sql)
            # 获取sql查询的列名
            desc = [item[0] for item in cursor.description]
            # 获取sql查询的结果
            data = cursor.fetchone()
            return dict(zip(desc, data)) if data else dict()

    def query_many(self, conn, sql):
        '''
        查询所有数据
        :param conn:数据库连接
        :param sql:查询sql
        :return:list
        '''

        # 获取数据库游标
        with conn.cursor() as cursor:
            # 执行查询sql语句
            cursor.execute(sql)
            # 获取sql查询的列名
            desc = [item[0] for item in cursor.description]
            # 获取sql查询的结果
            data = cursor.fetchall()
            return [dict(zip(desc, item)) for item in data] if data else list()

    def execute_one(self, conn=None, sql=None):
        '''
        执行单条 插入、更新、删除sql
        :param conn:数据库连接
        :param sql:插入、更新、删除sql
        :return:
        '''
        # 获取数据库游标
        with conn.cursor() as cursor:
            try:
                cursor.execute(sql)
            except pymysql.Error as e:
                # 执行插入、更新、删除sql时发生异常，事务回滚
                conn.rollback()
                logger.error('Mysql数据库执行SQL异常，原因：{}'.format(str(e)))
            else:
                # 成功执行插入、更新、删除sql后，提交事务
                conn.commit()

    def execute_many(self, conn=None, sql=None, datas=None):
        '''
        执行单条 插入、更新、删除sql
        :param conn:数据库连接
        :param sql:插入、更新、删除sql
        :param datas:插入、更新、删除sql的参数，参数格式为元组列表，示例：[(1,2,3,4),(1,2,3,4)]
        :return:
        '''
        # 获取数据库游标
        with conn.cursor() as cursor:
            try:
                cursor.executemany(sql, datas)
            except pymysql.Error as e:
                # 执行插入、更新、删除sql时发生异常，事务回滚
                conn.rollback()
                logger.error('Mysql数据库执行SQL异常，原因：{}'.format(str(e)))
            else:
                # 成功执行插入、更新、删除sql后，提交事务
                conn.commit()

    def change_db(self, conn=None, db=''):
        '''
        切换数据库
        :param conn:数据库连接
        :param db:待切换的数据库
        :return:
        '''
        # 获取数据库游标
        with conn.cursor() as cursor:
            try:
                sql = 'use {db};'.format(db=db)
                cursor.execute(sql)
            except pymysql.Error as e:
                conn.rollback()
                logger.error('Mysql数据库执行SQL异常，原因：{}'.format(str(e)))
            else:
                # 成功执行sql后，提交事务
                conn.commit()


class MysqlSshClient(object):
    '''
    Mysql SSH客户端
    '''

    def __init__(self, host=None, port=3306, username=None, password=None, db=None, charset='utf8', ssh_host=None,
                 ssh_username=None, ssh_password=None):
        '''
        数据库工具类初始化方法
        :param host: 数据库地址
        :param port: 数据库商品
        :param username: 数据库用户名
        :param password: 数据库密码
        :param db: 数据库名
        :param charset: 数据库编码
        :param ssh_host: ssh地址
        :param ssh_username: ssh用户名
        :param ssh_password: ssh密码
        '''
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.db = db
        self.charset = charset
        self.ssh_host = ssh_host
        self.ssh_username = ssh_username
        self.ssh_password = ssh_password
        try:
            self.tunnel = SSHTunnelForwarder(
                ssh_address_or_host=(self.ssh_host, 22),
                ssh_username=self.ssh_username,
                ssh_password=self.ssh_password,
                remote_bind_address=(self.host, self.port))
            self.tunnel.start()

            self.db_pool = PooledDB(pymysql, mincached=1, maxcached=5, host='127.0.0.1',
                                    port=self.tunnel.local_bind_port, user=self.username, passwd=self.password,
                                    db=self.db, charset=self.charset)
        except Exception as e:
            self.tunnel.close()
            logger.error('Mysql数据库工具初始化异常，原因：{}'.format(str(e)))

    def get_conn(self):
        '''
        获取数据库连接
        :return:
        '''
        try:
            return self.db_pool.connection()
        except Exception as e:
            logger.error('获取Mysql数据库连接异常，原因：{}'.format(str(e)))

    def close_conn(self, conn):
        '''
        关闭数据库连接
        :param conn: 数据库连接
        :return:
        '''
        if conn:
            try:
                conn.close()
            except pymysql.Error as e:
                logger.error('关闭Mysql数据库连接异常，原因：{}'.format(str(e)))

    def close(self):
        '''
        断开连接
        :return:
        '''
        try:
            self.db_pool.close()
        except Exception as e:
            logger.error('Mysql数据库断开连接异常，原因：{}'.format(str(e)))
        finally:
            self.tunnel.close()

    def query_one(self, conn=None, sql=None):
        '''
        查询一条数据
        :param conn:数据库连接
        :param sql:查询sql
        :return:
        '''
        # 获取数据库游标
        with conn.cursor() as cursor:
            # 执行查询sql语句
            cursor.execute(sql)
            # 获取sql查询的列名
            desc = [item[0] for item in cursor.description]
            # 获取sql查询的结果
            data = cursor.fetchone()
            return dict(zip(desc, data)) if data else dict()

    def query_many(self, conn=None, sql=None):
        '''
        查询所有数据
        :param conn:数据库连接
        :param sql:查询sql
        :return:list
        '''
        # 获取数据库游标
        with conn.cursor() as cursor:
            # 执行查询sql语句
            cursor.execute(sql)
            # 获取sql查询的列名
            desc = [item[0] for item in cursor.description]
            # 获取sql查询的结果
            data = cursor.fetchall()
            return [dict(zip(desc, item)) for item in data] if data else list()

    def execute_one(self, conn=None, sql=None):
        '''
        执行单条 插入、更新、删除sql
        :param conn:数据库连接
        :param sql:插入、更新、删除sql
        :return:
        '''
        # 获取数据库游标
        with conn.cursor() as cursor:
            try:
                cursor.execute(sql)
            except pymysql.Error as e:
                # 执行插入、更新、删除sql时发生异常，事务回滚
                conn.rollback()
                logger.error('Mysql数据库执行SQL异常，原因：{}'.format(str(e)))
            else:
                # 成功执行插入、更新、删除sql后，提交事务
                conn.commit()

    def execute_many(self, conn=None, sql=None, datas=None):
        '''
        执行单条 插入、更新、删除sql
        :param conn:数据库连接
        :param sql:插入、更新、删除sql
        :param datas:插入、更新、删除sql的参数，参数格式为元组列表，示例：[(1,2,3,4),(1,2,3,4)]
        :return:
        '''
        # 获取数据库游标
        with conn.cursor() as cursor:
            try:
                cursor.executemany(sql, datas)
            except pymysql.Error as e:
                # 执行插入、更新、删除sql时发生异常，事务回滚
                conn.rollback()
                logger.error('Mysql数据库执行SQL异常，原因：{}'.format(str(e)))
            else:
                # 成功执行插入、更新、删除sql后，提交事务
                conn.commit()

    def change_db(self, conn=None, db=''):
        '''
        切换数据库
        :param conn:数据库连接
        :param db:待切换的数据库
        :return:
        '''
        # 获取数据库游标
        with conn.cursor() as cursor:
            try:
                sql = 'use {db};'.format(db=db)
                cursor.execute(sql)
            except pymysql.Error as e:
                conn.rollback()
                logger.error('Mysql数据库执行SQL异常，原因：{}'.format(str(e)))
            else:
                # 成功执行sql后，提交事务
                conn.commit()
